#!/bin/usr/env python3

import numpy as np
from scipy.io import wavfile
from scipy.fft import fft, fftfreq
import matplotlib.pyplot as plt


def wav_to_beep_sequence(wav_path, frame_duration=50):
    """
    将WAV文件转换为beep频率-时长列表

    参数:
        wav_path: WAV文件路径
        frame_duration: 每帧时长(毫秒)

    返回:
        [(frequency, duration), ...] 列表
    """
    # 读取WAV文件
    sample_rate, data = wavfile.read(wav_path)

    # 如果是立体声，取单声道
    if len(data.shape) > 1:
        data = data.mean(axis=1)

    # 标准化到[-1, 1]
    data = data / np.max(np.abs(data))

    # 计算每帧样本数
    frame_size = int(sample_rate * frame_duration / 1000)
    num_frames = int(len(data) / frame_size)

    beep_sequence = []

    for i in range(num_frames):
        # 获取当前帧
        start = i * frame_size
        end = start + frame_size
        frame = data[start:end]

        # 应用汉宁窗减少频谱泄漏
        window = np.hanning(len(frame))
        frame = frame * window

        # 计算FFT
        n = len(frame)
        yf = fft(frame)
        xf = fftfreq(n, 1 / sample_rate)

        # 只取正频率
        positive_freq = xf[: n // 2]
        magnitudes = np.abs(yf[: n // 2])

        # 找到主导频率（忽略接近0的频率）
        min_freq = 20  # 忽略低于20Hz的频率
        valid_indices = np.where(positive_freq > min_freq)
        if len(valid_indices[0]) == 0:
            dominant_freq = 0
        else:
            dominant_freq = positive_freq[valid_indices][
                np.argmax(magnitudes[valid_indices])
            ]

        # 添加到序列
        if len(beep_sequence) == 0:
            beep_sequence.append((dominant_freq, frame_duration))
        else:
            last_freq, last_duration = beep_sequence[-1]
            if abs(dominant_freq - last_freq) < 10:  # 频率相近则合并
                beep_sequence[-1] = (last_freq, last_duration + frame_duration)
            else:
                beep_sequence.append((dominant_freq, frame_duration))

    # 过滤掉0频率（静音段）
    beep_sequence = [(f, d) for f, d in beep_sequence if f > 0]

    return beep_sequence


# 使用示例
beep_sequence = wav_to_beep_sequence("assets/dji_startup.wav")
print(beep_sequence)

# 可以保存到文件
with open("assets/dji_startup.frequency.txt", "w") as f:
    for freq, duration in beep_sequence:
        f.write(f"{freq:.1f},{duration}\n")
